from selenium import webdriver
from selenium.webdriver.edge.service import Service
from selenium.webdriver.common.by import By
from selenium.webdriver.common.keys import Keys
from selenium.webdriver.edge.options import Options
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as EC
from selenium.common.exceptions import TimeoutException
import requests  # Ensure requests is imported at the beginning of your file
import os
import time
import config
import re


def setup_driver():
    # Setup Edge options
    edge_options = Options()
    edge_options.add_argument("--enable-chrome-browser-cloud-management")
    edge_options.add_experimental_option("prefs", {
        "download.default_directory": config.DOWNLOAD_DIR,
        "download.prompt_for_download": False,
        "download.directory_upgrade": True,
        "safebrowsing.enabled": True
    })

    # Setup WebDriver
    service = Service(config.EDGE_DRIVER_PATH)
    driver = webdriver.Edge(service=service, options=edge_options)
    return driver


last_downloaded_audio_url = None
def download_audio(driver, text, line_number):
    global last_downloaded_audio_url  # 声明全局变量
    try:
        # 定位特定的<textarea>元素
        textarea_element = driver.find_element(By.XPATH, "//textarea[@data-testid='textbox']")
        if "本WebUI使用完全免费" in textarea_element.get_attribute('placeholder'):
            textarea_element.clear()  # 清除已有内容
            textarea_element.send_keys(text)  # 输入text
        if "我梦见一片焦土" in textarea_element.get_attribute('placeholder'):
            textarea_element.clear()  # 清除已有内容
            textarea_element.send_keys(text)  # 输入text

        # Find and print all buttons on the page
        buttons = driver.find_elements(By.TAG_NAME, "button")
        for button in buttons:
            if button.text == "切分生成":
                button.click()
                print("Clicked on '切分生成' button.")
                break
            if button.text == "生成音频！":
                button.click()
                print("Clicked on '生成音频' button.")
                break
            
        # 使用WebDriverWait等待音频链接出现
        try:
            # 等待音频链接出现
            audio_link = WebDriverWait(driver, 10).until(
                EC.visibility_of_element_located((By.XPATH, "//a[contains(@href, '.wav')]"))
            )
            audio_url = audio_link.get_attribute('href')
                        # Initialize a timeout for waiting for a new URL
            timeout = time.time() + 60*2  # 2 minutes from now
            while audio_url == last_downloaded_audio_url:
                if time.time() > timeout:
                    print("Timeout reached without new audio URL.")
                    break
                print("Waiting for new audio URL...")
                time.sleep(1)  # wait for 5 seconds before checking again
                audio_link = WebDriverWait(driver, 10).until(
                    EC.visibility_of_element_located((By.XPATH, "//a[contains(@href, '.wav')]"))
                )
                audio_url = audio_link.get_attribute('href')

            if audio_url != last_downloaded_audio_url:
                print("Audio URL:", audio_url)
                response = requests.get(audio_url)
                if response.status_code == 200:
                    file_path = os.path.join(config.DOWNLOAD_DIR, f"downloaded_audio_{line_number}.wav")
                    with open(file_path, "wb") as file:
                        file.write(response.content)
                    print(f"Audio file downloaded successfully: {file_path}")
                    # Update the last downloaded audio URL
                    last_downloaded_audio_url = audio_url
                else:
                    print("Failed to download the audio file.")
    
        except TimeoutException:
            print("Timeout waiting for the audio URL to appear.")

    except Exception as e:
        print(f"An error occurred: {e}")

def main():
    if not os.path.exists(config.DOWNLOAD_DIR):
        os.makedirs(config.DOWNLOAD_DIR)
        print(f"Created directory: {config.DOWNLOAD_DIR}")
    # Setup the WebDriver
    driver = setup_driver()
    # Navigate to the website
    driver.get(config.WEBSITE_URL)
    user_input = input("Please enter 'begin' to start the program: ")
    while user_input.lower() != "begin":
        user_input = input("Incorrect input. Please enter 'begin' to start: ")
    # Open input.txt and read lines
    with open("input.txt", "r", encoding="utf-8") as file:
        for line_number, text in enumerate(file, start=1):
            text = text.strip()  # Remove any leading/trailing whitespace
            if text:  # Check if the line is not empty
                download_audio(driver, text, line_number)

if __name__ == "__main__":
    main()
